package xyz.thoughtset.viewer.executor.blocks.executor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import xyz.thoughtset.viewer.ai.mcp.client.factory.McpBotFactory;
import xyz.thoughtset.viewer.executor.blocks.tool.FunctionCallBlockToolCallbackProvider;
import xyz.thoughtset.viewer.modules.step.entity.BlockTypeEnum;
import xyz.thoughtset.viewer.modules.step.entity.FunctionCallBody;
import xyz.thoughtset.viewer.modules.step.entity.ToolsBody;
import xyz.thoughtset.viewer.modules.step.entity.block.BlockBodyEle;
import xyz.thoughtset.viewer.modules.step.entity.block.BlockInfo;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;


@Slf4j
@Component
public class ToolsBlockExecutor extends AIChatExecutor<ToolsBody> {
    @Autowired
    private FunctionCallBlockToolCallbackProvider provider;
    @Autowired
    private McpBotFactory botFactory;

    @Override
    BlockTypeEnum getSupportType() {
        return BlockTypeEnum.TOOLS;
    }

    @Override
    protected void loadToolCallbacks(BlockInfo block,ToolsBody body, ChatClient.Builder builder) {
        List<BlockBodyEle> funs = body.getBodyEles();
        List<ToolCallback> toolCallbackList = new ArrayList<>();
        if (!ObjectUtils.isEmpty(funs)){
            Map<String,BlockBodyEle> funParamMap = funs.parallelStream().collect(Collectors.toMap(
                    BlockBodyEle::getEleId, Function.identity()
            ));
            toolCallbackList.addAll(
                    Arrays.asList(provider.getToolCallbacks(block.getId(),funParamMap))
            );
        }
        String clientIds = body.getMcpClientIds();
        if (StringUtils.hasText(clientIds)){
            toolCallbackList.addAll(
                    Arrays.asList(botFactory.clientsTools(
                            List.of(clientIds.split(","))
                    ))
            );
        }
        builder.defaultToolCallbacks(toolCallbackList);
    }


}
